# -*- coding: UTF-8 -*-
# __author__ = 'Sengo'

import cv2
import numpy as np


def ECC(im1, im2, original1=None, original2=None, i=0):
    # Convert images to grayscale
    im1_gray = im1
    im2_gray = im2

    # Find size of image1
    if None is not original1:
        sz = original1.shape
    else:
        original2 = im2
        sz = im1.shape
    # Define the motion model
    warp_mode = cv2.MOTION_TRANSLATION

    # Define 2x3 or 3x3 matrices and initialize the matrix to identity
    if warp_mode == cv2.MOTION_HOMOGRAPHY:
        warp_matrix = np.eye(3, 3, dtype=np.float32)
    else:
        warp_matrix = np.eye(2, 3, dtype=np.float32)

    # Specify the number of iterations.
    number_of_iterations = 5000

    # Specify the threshold of the increment
    # in the correlation coefficient between two iterations
    termination_eps = 1e-10

    # Define termination criteria
    criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, number_of_iterations, termination_eps)

    # Run the ECC algorithm. The results are stored in warp_matrix.
    (cc, warp_matrix) = cv2.findTransformECC(im1_gray, im2_gray, warp_matrix, warp_mode, criteria)

    # 根据图像金字塔缩小了几次就乘以几次2倍的矩阵

    while 1:
        if i == 0:
            break
        m = np.mat([[1, 1, 2], [1, 1, 2]])
        warp_matrix = np.multiply(warp_matrix, m)
        i -= 1

    if warp_mode == cv2.MOTION_HOMOGRAPHY:
        # Use warpPerspective for Homography
        im2_aligned = cv2.warpPerspective(im2, warp_matrix, (sz[1], sz[0]),
                                          flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP)
    else:
        # print warp_matrix
        # Use warpAffine for Translation, Euclidean and Affine
        im2_aligned = cv2.warpAffine(original2, warp_matrix, (sz[1], sz[0]),
                                     flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP)

    # Show final results
    # cv2.imshow("Image 1", im1)
    # cv2.imshow("Image 2", im2)
    # cv2.imshow("Aligned Image 2", im2_aligned)
    # cv2.waitKey(0)
    return im2_aligned
